from datetime import datetime, date
import backtrader as bt
import akshare as ak
import pandas as pd
import numpy as np
import quantstats as qs
import pyfolio as pf
import webbrowser
from sqlalchemy import create_engine, distinct, or_, and_
import pymssql
from urllib.parse import quote_plus as urlquote
from configparser import ConfigParser
from sklearn.linear_model import LinearRegression
import warnings

warnings.filterwarnings('ignore')

pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)

conf = ConfigParser()
conf.read('env.ini')

sqlserver = ('sqlserver_out', 'sqlserver_in')[0]

host = conf.get(sqlserver, 'host')
port = conf.get(sqlserver, 'port')
user = conf.get(sqlserver, 'user')
password = conf.get(sqlserver, 'password')
schema = conf.get(sqlserver, 'schema')
db_name = conf.get(sqlserver, 'db')


def get_stocks():
    sql = "SELECT STK_CODE FROM ETF_INDEX_INFO"
    r = exec_sql(sql)
    df1 = pd.DataFrame(r, columns=['代码'])
    df2 = ak.fund_etf_spot_em()
    res = pd.merge(df1, df2, on=['代码'])
    res = res.loc[res['最新价'] < 10, :]
    stocks = res['代码'].to_list()
    # stocks = [tok for tok in stocks if tok != '512870']

    return stocks


def get_data(stocks, freq='day', source='db'):
    if freq == '1min':
        table_name = 'STOCK_MIN_DATA'
    elif freq == '15min':
        table_name = 'STOCK_15MIN_DATA'
    else:
        table_name = 'ETF_DAY_DATA'  # HIS_DAILY_DATA

    res = dict()
    date_data = set_date()
    st = date_data['st']
    et = date_data['et']
    st_d = date_data['st_d']
    et_d = date_data['et_d']

    if source == 'db':
        print('db')
        stocks_str = ','.join(f"'{tok}'" for tok in stocks)
        sql = f"SELECT code, datetime, [close], [open], low, high, volume FROM {table_name} WHERE code in ({stocks_str})"
        r = exec_sql(sql)
        df = pd.DataFrame(r, columns=['code', 'datetime', 'close', 'open', 'low', 'high', 'volume'])
        # df['datetime'] = df['datetime'].map(time_map)

        for code in stocks:
            if code == '399300':
                sig_df = ak.stock_zh_index_daily(symbol="sz399300")
                sig_df = sig_df.loc[(sig_df.date >= st_d) & (sig_df.date <= et_d), :]
                sig_df.rename(columns={'date': 'datetime'}, inplace=True)
                sig_df.insert(0, 'code', code)
                sig_df = sig_df.loc[:, ['code', 'datetime', 'close', 'open', 'low', 'high', 'volume']]
                # sig_df['datetime'] = pd.to_datetime(sig_df['datetime'])
            else:
                sig_df = df.loc[df['code'] == code]
            sig_df['datetime'] = pd.to_datetime(sig_df['datetime'])
            sig_df.sort_values('datetime', inplace=True, ignore_index=True)
            sig_df.set_index('datetime', inplace=True)
            print(f'code={code}->len={len(sig_df)}')
            res[code] = sig_df
    else:
        print('internet')
        for code in stocks:
            if code == '399300':
                df = ak.stock_zh_index_daily(symbol="sz399300")
                df = df.loc[(df.date >= st_d) & (df.date <= et_d), :]
                df.rename(columns={'date': 'datetime'}, inplace=True)
            else:
                df = ak.fund_etf_hist_em(symbol=code, period="daily", start_date=st, end_date=et, adjust="qfq")
                df.rename(columns={'日期': 'datetime', '开盘': 'open', '收盘': 'close', '最高': 'high',
                                   '最低': 'low', '成交量': 'volume'}, inplace=True)

            df.insert(0, 'code', code)
            df = df.loc[:, ['code', 'datetime', 'close', 'open', 'low', 'high', 'volume']]
            df['datetime'] = pd.to_datetime(df['datetime'])
            df.set_index('datetime', inplace=True)
            print(f'code={code}->len={len(df)}')
            res[code] = df

    return res


def exec_sql(sql):
    conn = pymssql.connect(host=host, port=port, user=user, password=password, database=db_name)
    cursor = conn.cursor()
    cursor.execute(sql)
    r = cursor.fetchall()
    cursor.close()
    conn.close()
    return r


def time_map(x):
    t = datetime.strptime(str(x), '%Y-%m-%d')
    # t = datetime.strptime(str(x), '%Y%m%d%H%M')
    return t


def insert_db(df, table_name):
    yconnect = create_engine(f'mssql+pymssql://{user}:{urlquote(password)}@{host}:{port}/{db_name}?charset=utf8')
    pd.io.sql.to_sql(df, table_name, yconnect, schema=schema, if_exists='append', index=False)


pre_ratio = 0.0048
ratio = 0.1
trail_ratio = 0.035
count_info = dict()
# stocks = get_stocks()
init_cash = 1e7


def run_backtrade(stocks):
    cerebro = bt.Cerebro()
    cerebro.addstrategy(GridReversionStrategy)
    data_all = get_data(stocks, source='db')
    index = stocks[0]
    data0 = data_all[index]
    datafeed0 = bt.feeds.PandasData(dataname=data0)
    cerebro.adddata(datafeed0, name=f'{index}')

    for i, code in enumerate(data_all):
        if code == index:
            continue
        print(f"i={i}->{code} Done !")
        data = data_all[code]

        datafeed = bt.feeds.PandasData(dataname=data)
        # st_date = datetime(2023, 1, 1)
        # ed_date = datetime(2023, 6, 16)
        # datafeed = bt.feeds.PandasData(dataname=data, fromdate=st_date, todate=ed_date)
        cerebro.adddata(datafeed, name=code)
        print(f'{code} feeds ok')

    # 初始资金 1,000,000.0
    cerebro.broker.setcash(init_cash)
    # 佣金，双边各 0.0003
    cerebro.broker.setcommission(commission=0.0003)
    # 滑点：双边各 0.0001
    cerebro.broker.set_slippage_perc(perc=0.0001)
    # cerebro.addsizer(bt.sizers.FixedSize, stake=100)

    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='pnl')  # 返回收益率时序数据
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name='_AnnualReturn')  # 年化收益率
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='_SharpeRatio')  # 夏普比率
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='_DrawDown')  # 回撤
    cerebro.addanalyzer(bt.analyzers.PyFolio, _name='pyfolio')

    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    result = cerebro.run()
    strat = result[0]
    st = strat.params.st
    pyfoliozer = strat.analyzers.getbyname('pyfolio')
    # 返回日度收益率序列
    daily_return = pd.Series(strat.analyzers.pnl.get_analysis())
    # 打印评价指标
    print("--------------- TimeReturn -----------------")
    print(daily_return)
    print("--------------- AnnualReturn -----------------")
    print(strat.analyzers._AnnualReturn.get_analysis())
    print("--------------- SharpeRatio -----------------")
    print(strat.analyzers._SharpeRatio.get_analysis())
    print("--------------- DrawDown -----------------")
    print(strat.analyzers._DrawDown.get_analysis())
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())

    returns, positions, transactions, gross_lev = pyfoliozer.get_pf_items()
    print(st.date())
    returns = returns[returns.index.date >= st.date()]
    print(f'returns={returns}')
    html_file = 'br.html'
    qs.reports.html(returns, title='grid', output=html_file, download_filename=html_file)
    qs.reports.basic(returns)
    f = webbrowser.open(html_file)
    # pf.create_full_tear_sheet(returns)
    # cerebro.plot()


class GridReversionStrategy(bt.Strategy):
    params = (
        ('period', 25),
        ('init_money', 3e4),
        ('add_money', 1e4),
        ('buy_ratio', -0.35),
        ('add_ratio', -0.1),
        ('sell_ratio', 0.2),
        ('printlog', False),
        ('st', datetime(2013, 8, 1)),
    )

    def __init__(self):
        self.order = None
        self.daily_pnl = dict()
        self.trade_info = list()
        self.pos_info = dict()
        self.trail_dict = dict()
        self.buy_first = dict()  # 每支ETF是否开始第一天就跌破之前最高价的buy_ratio
        self.TR = dict()
        self.ATR = dict()

        stocks = list()
        for data in self.datas:
            code = data._name
            if code != '399300':
                stocks.append(code)
                self.trail_dict[code] = {'hdt': 0, 'high': 0, 'ldt': 0, 'low': 0, 'cdt': 0, 'cost': 0, 'buy_count': 0,
                                         'trend': 'x', 'pos': dict(), 'grid': dict(), 'trail_high': False}
                self.buy_first[code] = True

        # self.trade_date = get_trade_date(stocks)

    def downcast(self, price, style='add'):
        res = self.p.add_money / price // 100 * 100 if style == 'add' else self.p.init_money / price // 100 * 100
        return res

    def set_grid(self, code, price):
        self.trail_dict[code]['grid']['lev0'] = price
        self.trail_dict[code]['grid']['lev_m'] = price * (1 + 2.7*self.p.add_ratio)
        self.trail_dict[code]['grid']['lev1'] = price * (1 + self.p.buy_ratio)
        self.trail_dict[code]['grid']['lev2'] = self.trail_dict[code]['grid']['lev1'] * (1 + self.p.add_ratio)
        self.trail_dict[code]['grid']['lev3'] = self.trail_dict[code]['grid']['lev2'] * (1 + self.p.add_ratio)
        self.trail_dict[code]['grid']['lev4'] = self.trail_dict[code]['grid']['lev3'] * (1 + self.p.add_ratio)
        self.log(f"high={price}->lev_m={price * (1 + self.p.add_ratio)}->lev1={price * (1 + self.p.buy_ratio)}->"
                 f"lev2={self.trail_dict[code]['grid']['lev1'] * (1 + self.p.add_ratio)}->"
                 f"lev3={self.trail_dict[code]['grid']['lev2'] * (1 + self.p.add_ratio)}->"
                 f"lev4={self.trail_dict[code]['grid']['lev3'] * (1 + self.p.add_ratio)}")
        self.log(50 * '-')

    def prenext(self):
        self.next()

    def log(self, txt, dt=None, doprint=True):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            print('%s, %s' % (dt.isoformat(), txt))

    @staticmethod
    def percent(today, yesterday):
        return float(today - yesterday) / today

    def next(self):
        if self.order:
            return

        cur_dt = self.datas[0].datetime.datetime(0)

        if cur_dt < self.p.st:
            for data in self.datas:
                code = data._name
                try:
                    dt = data.datetime.datetime(0)  # 获取当前的回测时间点
                except:
                    print(f'{code}-{dt}')
                if code == '399300' or dt != cur_dt:
                    continue
                if data.high[0] > self.trail_dict[code]['high']:
                    self.trail_dict[code]['hdt'] = dt
                    self.trail_dict[code]['high'] = data.high[0]
            return

        dt_int = int(cur_dt.strftime('%Y%m%d'))

        if dt_int not in self.daily_pnl:
            self.daily_pnl[dt_int] = round(self.broker.getvalue() - init_cash, 2)

        for data in self.datas:
            code = data._name
            try:
               dt = data.datetime.datetime(0)  # 获取当前的回测时间点
            except:
                print(f'{code}-{dt}')
            # dt != cur_dt 防止停牌日触发买卖
            if code == '399300' or dt != cur_dt:
                continue

            # self.log(f"first={self.buy_first[code]}")

            if not self.buy_first[code] and data.high[0] > self.trail_dict[code]['high']:
                self.trail_dict[code]['hdt'] = dt
                self.trail_dict[code]['high'] = data.high[0]
                self.set_grid(code, data.high[0])
                self.trail_dict[code]['trail_high'] = True
                continue
            elif self.buy_first[code]:
                self.set_grid(code, self.trail_dict[code]['high'])

            if self.buy_first[code]:
                if data.low[0] < self.trail_dict[code]['grid']['lev4']:
                    # amount1 = self.downcast(data.close[0], style='init')
                    amount1 = self.downcast(data.low[0])
                    amount2 = self.downcast(data.low[0], style='init')
                    amount3 = 3 * amount2
                    amount4 = 6 * amount2
                    amount = amount1 + amount2 + amount3 + amount4
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['trend'] = 'down'
                    self.trail_dict[code]['buy_count'] = 4
                    self.trail_dict[code]['pos'][1] = amount1
                    self.trail_dict[code]['pos'][2] = amount2
                    self.trail_dict[code]['pos'][3] = amount3
                    self.trail_dict[code]['pos'][3] = amount4
                    self.log(f'buy lev4->a1={amount1}->a2={amount2}->a3={amount3}->a4={amount4}')
                elif data.low[0] < self.trail_dict[code]['grid']['lev3']:
                    # amount1 = self.downcast(data.close[0], style='init')
                    amount1 = self.downcast(data.low[0])
                    amount2 = self.downcast(data.low[0], style='init')
                    amount3 = 3 * amount2
                    amount = amount1 + amount2 + amount3
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['trend'] = 'down'
                    self.trail_dict[code]['buy_count'] = 3
                    self.trail_dict[code]['pos'][1] = amount1
                    self.trail_dict[code]['pos'][2] = amount2
                    self.trail_dict[code]['pos'][3] = amount3
                    self.log(f'buy lev3->a1={amount1}->a2={amount2}->a3={amount3}')
                elif data.low[0] < self.trail_dict[code]['grid']['lev2']:
                    # amount1 = self.downcast(data.close[0], style='init')
                    amount1 = self.downcast(data.low[0])
                    amount2 = self.downcast(data.low[0], style='init')
                    amount = amount1 + amount2
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['trend'] = 'down'
                    self.trail_dict[code]['buy_count'] = 2
                    self.trail_dict[code]['pos'][1] = amount1
                    self.trail_dict[code]['pos'][2] = amount2
                    self.log(f'buy lev2->a1={amount1}->a2={amount2}')
                elif data.low[0] < self.trail_dict[code]['grid']['lev1']:
                    # amount = self.downcast(data.close[0], style='init')
                    amount = self.downcast(data.low[0])
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['trend'] = 'down'
                    self.trail_dict[code]['buy_count'] = 1
                    self.trail_dict[code]['pos'][1] = amount
                    self.log(f'buy lev1->a={amount}')
                else:
                    self.buy_first[code] = False
            else:
                if not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 0 and data.low[0] < \
                        self.trail_dict[code]['grid']['lev1']:
                    amount = self.downcast(data.low[0])
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['trend'] = 'down'
                    self.log(f'buy lev1->xx_a={amount}')
                    # self.trail_dict[code]['buy_count'] += 1
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 1 and data.low[
                    0] < self.trail_dict[code]['grid']['lev2']:
                    amount = self.downcast(data.low[0], style='init')
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['trend'] = 'down'
                    self.log(f'buy lev2->xx_a={amount}')
                    # self.trail_dict[code]['buy_count'] += 1
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 2 and data.low[
                    0] < self.trail_dict[code]['grid']['lev3']:
                    amount = 3 * self.downcast(data.low[0], style='init')
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['trend'] = 'down'
                    self.log(f'buy lev3->xx_a={amount}')
                    # self.trail_dict[code]['buy_count'] += 1
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 3 and data.low[
                    0] < self.trail_dict[code]['grid']['lev4']:
                    amount = 6 * self.downcast(data.low[0], style='init')
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['trend'] = 'down'
                    self.log(f'buy lev4->xx_a={amount}')
                    # self.trail_dict[code]['buy_count'] += 1

                # 卖出
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 4 and data.high[
                    0] > self.trail_dict[code]['grid']['lev2']:
                    amount = self.trail_dict[code]['pos'][4]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_1'
                    self.log(f'sell lev2->in_a={amount}')
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 3 and data.high[
                    0] > self.trail_dict[code]['grid']['lev1']:
                    amount = self.trail_dict[code]['pos'][3]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_2'
                    self.log(f'sell lev1->in_a={amount}')
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 2 and data.high[
                    0] > self.trail_dict[code]['grid']['lev_m']:
                    amount = self.trail_dict[code]['pos'][2]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_3'
                    self.log(f'sell lev_m->in_a={amount}')
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 1 and data.high[
                    0] > self.trail_dict[code]['grid']['lev0']*0.85:
                    amount = self.trail_dict[code]['pos'][1]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_4'
                    self.log(f'sell lev0->in_a={amount}')
                elif self.trail_dict[code]['trail_high'] and data.close[0] < self.trail_dict[code]['high'] * 0.9:
                    self.trail_dict[code]['trail_high'] = False
                    self.order = self.close(data=data, name=code)

                if data.low[0] < self.trail_dict[code]['low']:
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]

    def notify_order(self, order):
        code = order.data._name

        if order.status in [order.Submitted, order.Accepted]:
            return

        if order.status in [order.Completed]:

            if order.isbuy():
                self.log(
                    f'code={code} BUY CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log(
                    f'code={code} BUY EXECUTED, ref:{order.ref}，Price: {order.executed.price:.2f}, Cost: {order.executed.value: .2f},'
                    f'Comm {order.executed.comm:.2f}, Size: {order.executed.size:.2f}')

                trade_dt = int(bt.num2date(order.created.dt).strftime('%Y%m%d'))
                trade_item = [trade_dt, code, 'BUY', round(order.executed.price, 2), order.executed.size,
                              round(order.executed.comm, 2)]
                self.trade_info.append(trade_item)
                if self.buy_first[code]:
                    self.buy_first[code] = False
                else:
                    self.trail_dict[code]['buy_count'] += 1
                    self.trail_dict[code]['pos'][self.trail_dict[code]['buy_count']] = order.executed.size
                self.trail_dict[code]['cdt'] = int(bt.num2date(order.executed.dt).strftime('%Y%m%d'))
                self.trail_dict[code]['cost'] = round(order.executed.price, 2)
            else:
                self.log(
                    f'code={order.data._name} SELL CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log(
                    f'code={order.data._name} SELL EXECUTED, ref:{order.ref}，Price: {order.executed.price:.2f}, Cost: {order.executed.value: .2f},'
                    f' Comm {order.executed.comm:.2f}, Size: {order.executed.size:.2f}')
                trade_dt = int(bt.num2date(order.created.dt).strftime('%Y%m%d'))
                trade_item = [trade_dt, code, 'SELL', round(order.executed.price, 2),
                              order.executed.size, round(order.executed.comm, 2)]
                self.trade_info.append(trade_item)
                self.trail_dict[code]['cdt'] = int(bt.num2date(order.executed.dt).strftime('%Y%m%d'))
                self.trail_dict[code]['cost'] = round(order.executed.price, 2)
                del self.trail_dict[code]['pos'][self.trail_dict[code]['buy_count']]
                self.trail_dict[code]['buy_count'] -= 1
            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log(f'Order Canceled/Margin/Rejected->code={code}')

        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        self.log(f'OPERATION PROFIT, GROSS {trade.pnl:.2f}, NET {trade.pnlcomm:.2f}')

    def stop(self):
        print('Strategy Finish!!!')

        # with open(r'C:\Users\AndrewX\Desktop\cache_log.txt', 'w') as f:
        #     res = '\n'.join(self.cache_log)
        #     f.write(res)
        doprint = (False, True)[0]

        if doprint:
            daily_pnl_df = pd.DataFrame(self.daily_pnl.items(), columns=['DATE_T', 'PNL'])
            append_df = pd.DataFrame([[101, round((self.broker.getvalue() / init_cash - 1) * 100, 2)]],
                                     columns=['DATE_T', 'PNL'])
            daily_pnl_df = daily_pnl_df.append(append_df, ignore_index=True)
            trade_df = pd.DataFrame(self.trade_info,
                                    columns=['DATE_T', 'STK_CODE', 'ACTION', 'PRICE', 'AMOUNT', 'COMM'])
            pnl_table = 'BREAKUPSTOCK_DAILY_PNL'
            trade_table = 'BREAKUPSTOCK_TRADE_INFO'
            # daily_pnl_df = daily_pnl_df.loc[(daily_pnl_df['DATE_T'] >= 20230428) | (daily_pnl_df['DATE_T'] == 101)]
            insert_db(daily_pnl_df, pnl_table)
            insert_db(trade_df, trade_table)
            print(self.trail_dict)
            print('insert table ok')


# 计算强弱得分
def calculate_score(srs, period=25):
    if srs.shape[0] < period:
        return np.nan
    x = np.arange(1, period + 1)
    # print(f'xx={srs.values / srs.values[0]}->type={type(srs.values / srs.values[0])}')
    y = srs.values / srs.values[0]
    lr = LinearRegression().fit(x.reshape(-1, 1), y)
    # 斜率
    slope = lr.coef_[0]
    # 决定系数R2
    r_squared = lr.score(x.reshape(-1, 1), y)
    # 得分
    score = 10000 * slope * r_squared
    return score


def get_trade_date(stocks):
    n = 25
    date_data = set_date()
    st = date_data['st']
    et = date_data['et']
    st_d = date_data['st_d']
    et_d = date_data['et_d']
    ret = ak.stock_zh_index_daily(symbol="sz399300")
    ret = ret.loc[(ret['date'] >= st_d) & (ret['date'] <= et_d), ['date']]
    ret.rename(columns={'date': 'datetime'}, inplace=True)
    ret['datetime'] = pd.to_datetime(ret['datetime'])

    for i, code in enumerate(stocks):
        df = ak.fund_etf_hist_em(symbol=code, period="daily", start_date=st, end_date=et, adjust="qfq")
        df = df.loc[:, ['日期', '收盘']]
        df.rename(columns={'日期': 'datetime', '收盘': code}, inplace=True)
        df['datetime'] = pd.to_datetime(df['datetime'])
        ret = pd.merge(ret, df, on=['datetime'], how='left')
    # ret.set_index('datetime', inplace=True)
    ret.fillna(method='ffill', inplace=True)

    for code in stocks:
        ret[code + '_score'] = ret.loc[:, code].rolling(n).apply(lambda x: calculate_score(x, n))

    ret = ret.dropna()
    ret['signal'] = ret[[code + '_score' for code in stocks]].idxmax(axis=1).str.replace('_score', '')

    return ret.loc[:, ['datetime', 'signal']]


def set_date():
    st = "20130101"
    et = "20230816"
    st_d = datetime.strptime(st, '%Y%m%d').date()
    et_d = datetime.strptime(et, '%Y%m%d').date()
    res = {'st': st, 'et': et, 'st_d': st_d, 'et_d': et_d}

    return res


def main():
    index = '399300'
    # stocks = get_stocks()
    # stocks = ['510880', '159915', '513100', '518880']
    # stocks = ['159901', '510880']
    # stocks = ['159901', '510880', '512510', '512220', '159938', '159909', '513500', '513880', '515770', '515350', '159810', '517350']
    # stocks = stocks[:200]
    # stocks = ['510880']
    # stocks = ['159909', '510170', '159915', '512220', '510880']
    stocks = ['510300', '517770']
    stocks.insert(0, index)
    run_backtrade(stocks)
    # s = get_data(stocks, freq='day', source='not db')
    # print(s)


if __name__ == '__main__':
    main()
